﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;

namespace Advanced.Algorithms.DataStructures;

/// <summary>
///     A disjoint set implementation.
/// </summary>
public class DisJointSet<T> : IEnumerable<T>
{
    /// <summary>
    ///     A Map for faster access for members.
    /// </summary>
    private readonly Dictionary<T, DisJointSetNode<T>> set = new();

    public int Count { get; private set; }

    IEnumerator IEnumerable.GetEnumerator()
    {
        return GetEnumerator();
    }

    public IEnumerator<T> GetEnumerator()
    {
        return set.Values.Select(x => x.Data).GetEnumerator();
    }

    /// <summary>
    ///     Creates a new set with given member.
    ///     Time complexity: log(n).
    /// </summary>
    public void MakeSet(T member)
    {
        if (set.ContainsKey(member)) throw new Exception("A set with given member already exists.");

        var newSet = new DisJointSetNode<T>
        {
            Data = member,
            Rank = 0
        };

        //Root's Parent is Root itself
        newSet.Parent = newSet;
        set.Add(member, newSet);

        Count++;
    }


    /// <summary>
    ///     Returns the reference member of the set where this member is part of.
    ///     Time complexity: log(n).
    /// </summary>
    public T FindSet(T member)
    {
        if (!set.ContainsKey(member)) throw new Exception("No such set with given member.");

        return FindSet(set[member]).Data;
    }

    /// <summary>
    ///     Recursively move up in the set tree till root
    ///     and return the Root.
    ///     Does path Compression on all visited members on way to root
    ///     by pointing their parent to Root.
    /// </summary>
    private DisJointSetNode<T> FindSet(DisJointSetNode<T> node)
    {
        var parent = node.Parent;

        if (node != parent)
        {
            //compress path by setting parent to Root
            node.Parent = FindSet(node.Parent);
            return node.Parent;
        }

        //reached root so return the Root (reference Member)
        return parent;
    }

    /// <summary>
    ///     Union's given member's sets if given members are in differant sets.
    ///     Otherwise does nothing.
    ///     Time complexity: log(n).
    /// </summary>
    public void Union(T memberA, T memberB)
    {
        var rootA = FindSet(memberA);
        var rootB = FindSet(memberB);

        if (rootA.Equals(rootB)) return;

        var nodeA = set[rootA];
        var nodeB = set[rootB];

        //equal rank so just pick any of two as Root
        //and increment rank
        if (nodeA.Rank == nodeB.Rank)
        {
            nodeB.Parent = nodeA;
            nodeA.Rank++;
        }
        else
        {
            //pick max Rank node as root
            if (nodeA.Rank < nodeB.Rank)
                nodeA.Parent = nodeB;
            else
                nodeB.Parent = nodeA;
        }
    }
}

internal class DisJointSetNode<T>
{
    internal T Data { get; set; }
    internal int Rank { get; set; }

    internal DisJointSetNode<T> Parent { get; set; }
}